#!/usr/bin/env pwsh

$ErrorActionPreference = "Stop"
Set-StrictMode -Version Latest

# 参数处理
$CONFIG = $args[0]
$GPUS = $args[1]

# 设置环境变量默认值
if (-not $env:NNODES) { $env:NNODES = 1 }
if (-not $env:NODE_RANK) { $env:NODE_RANK = 0 }
if (-not $env:PORT) { $env:PORT = 29500 }
if (-not $env:MASTER_ADDR) { $env:MASTER_ADDR = "127.0.0.1" }

# 关键修复：强制设置 USE_LIBUV=0
$env:USE_LIBUV = "0"

# 设置PYTHONPATH
$scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Definition
$pythonPath = Join-Path $scriptDir ".."
if ($env:PYTHONPATH) {
    $env:PYTHONPATH = "$pythonPath;$env:PYTHONPATH"
} else {
    $env:PYTHONPATH = $pythonPath
}

# 构建参数数组
$pyArgs = @(
    "--nnodes=$env:NNODES"
    "--node_rank=$env:NODE_RANK"
    "--master_addr=$env:MASTER_ADDR"
    "--nproc_per_node=$GPUS"
    "--master_port=$env:PORT"
    (Join-Path $scriptDir "train.py")
    $CONFIG
    "--launcher"
    "pytorch"
)

# 添加剩余参数（从第三个参数开始）
if ($args.Length -ge 3) {
    $pyArgs += $args[2..($args.Length-1)]
}

Write-Host "执行命令: torchrun $pyArgs"
torchrun @pyArgs